"""对比ode23、odeint、quad和solver_ivp"""
import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint, quad, solve_ivp
import scipy.io as scio

def noisePSD(freq, noise_power, Cp, tao1, tao2):
# def noisePSD(freq, Cp, tao1, tao2):
    """计算各种条件下的噪声PSD"""
    e_n, i_n = 20e-9, 40e-15
    Ceq = 2*(1.15e-12 + Cp)
    L, r = 5e-3, 31.415926
    Rf, Cf = 1e7, 3.3e-12

    lp_freq = abs(freq - 1e5)
    s = freq*1j*2*np.pi
    s_2 = -(freq*2*np.pi)**2
    my_Z_BR = (L*s + r)/(L*Ceq*s_2 + r*Ceq*s + 1)
    Z_f = Rf/(Rf*Cf*s + 1)    # 热噪声
    thermal_noise_BR = 0.1287*((my_Z_BR.real)**0.5)*1e-9# 桥路热电阻V/Hz^0.5
    thermal_noise_gain = -2*Z_f/my_Z_BR
    th_BR_Uo = thermal_noise_BR * abs(thermal_noise_gain )
    # 等效输入电压噪声
    noise_gain  = 1 - thermal_noise_gain 
    en_Uo = e_n * abs(noise_gain )
    # 等效输入电流噪声
    e_in_Uo = abs(Z_f )*i_n
    # 反馈阻抗热噪声
    th_Zf_Uo = 0.1287*((Z_f .real)**0.5)*1e-9
    # 总噪声
    before_psd = th_BR_Uo**2 + 2*(en_Uo**2 + e_in_Uo**2 + th_Zf_Uo**2)
    # (70Hz, 350HZ)滤波器作用
    temp = tao1*tao2
    lp_filter = temp/(-lp_freq**2 + (tao1 + tao2)*lp_freq*1j + temp)
    psd = before_psd * ((lp_filter.real)**2 + (lp_filter.imag)**2)
    return psd

tao1, tao2 = (70, 350)
Cp = 252.15e-12
matlab_data = scio.loadmat('./psd.mat')
delta_Cps = matlab_data['delta_Cps']
Freq_span = matlab_data['Freq_span']
Freq_span.shape = Freq_span.size,
matlab_psds = matlab_data['PSD'].T

# compare integrands from MATLAB and Python
psds = np.empty( (delta_Cps.size, Freq_span.size), dtype=np.float64)
index = 0
for delta_Cp in delta_Cps.flat:
    psds[index,:] = noisePSD(Freq_span, 0, \
        (delta_Cp+1)*Cp, tao1, tao2)
    index += 1
print(delta_Cps)
# 指定dpi=141.21后反而模糊，默认才80吧。。
fig,axes = plt.subplots(1, 2, figsize=(18/2.54, 9/2.54))
for index in range(0,3):
    axes[0].plot(Freq_span, psds[index, :], label=str(delta_Cps[0, index]))
    axes[1].plot(Freq_span, (psds[index,:]-matlab_psds[index,:])/\
        matlab_psds[index,:], label=str(delta_Cps[0, index]))
axes[1].legend(fontsize=10)
axes[1].xaxis.set_tick_params(labelsize=10)
axes[1].set_xlabel('frequency/Hz', fontsize=10)
axes[1].grid()
axes[0].legend(fontsize=10)
axes[0].xaxis.set_tick_params(labelsize=10)
axes[0].set_xlabel('frequency/Hz', fontsize=10)
axes[0].grid()
# fig.title('compare of integrand',fontsize=12)
plt.show()

# 对PSD进行积分，得到噪声功率
delta_Cps = np.arange(-1e-2, 1e-2+1e-4, 1e-4)
frequency = np.linspace(1e5-500, 1e5+500, 101)
df = frequency[1] - frequency[0]# freq 线性递增
noise_powers = np.empty((4, delta_Cps.size), dtype=np.float64)
index = 0
for delta_Cp in delta_Cps:
    # 使用odeint()
    noise_power = odeint(noisePSD, [0], frequency,\
        args=((delta_Cp+1)*Cp, tao1, tao2), tfirst=True)
    noise_powers[0, index] = noise_power[-1]
    # 使用quad()
    noise_power = quad(noisePSD, 1e5-500, 1e5+500,\
        args=(0, (delta_Cp+1)*Cp, tao1, tao2))# args是位置参数
    noise_powers[1, index] = noise_power[0]
    # 使用梯形法
    noise_psd = noisePSD(frequency, 0, (delta_Cp+1)*252.15e-12,\
        tao1, tao2)
    noise_power = 0.5*df*(noise_psd[:-1] + noise_psd[1:])
    noise_powers[2, index] = noise_power.sum()
    # 使用solver_ivp()
    result = solve_ivp(noisePSD, [1e5-500,1e5+500],np.array([0]), \
        method='RK23', t_eval=frequency, args=((delta_Cp+1)*Cp, tao1, tao2), \
        max_step=df)
    noise_powers[3, index] = result.y[0, -1]
    index += 1

fig, axes = plt.subplots(1, 2, sharex=True, figsize=(20/2.54, 10/2.54))
axes[1].plot(delta_Cps, noise_powers[0, :], label='odeint')
axes[1].legend(fontsize=10)
axes[1].set_xlabel(r'$\Delta Cp$', fontsize=10)
axes[1].xaxis.set_tick_params(labelsize=10)
axes[1].grid()
axes[0].plot(delta_Cps, noise_powers[1, :], label='quad')
axes[0].plot(delta_Cps, noise_powers[2, :], label='trapezoidal')
axes[0].plot(delta_Cps, noise_powers[3, :], label='solver_ivp')
axes[0].set_ylabel('noise power', fontsize=10)
axes[0].set_xlabel(r'$\Delta Cp$', fontsize=10)
axes[0].legend(fontsize=10)
axes[0].grid()
plt.show()

